#include "detector/dl/lattice_net_decode.h"

#include <numeric>

#include "detector/dl/cc_decode.h"
#include "detector/dl/lane_decode.h"

LatticeNetDecode::LatticeNetDecode() noexcept {
  obj_scores_ptr = std::make_shared<std::array<float, kMaxResultSize>>();
  obj_labels_ptr = std::make_shared<std::array<int32_t, kMaxResultSize>>();
  obj_boxes_ptr =
      std::make_shared<std::array<float, kMaxResultSize * kMaxObjInfoSize>>();
  obj_nms_keep_ptr = std::make_shared<std::array<bool, kMaxResultSize>>();

  map_pts_ptr_ = std::make_shared<
      std::array<float, kMaxResultSize * kMaxMapFeatureSize>>();
}
LatticeNetDecode::~LatticeNetDecode() {}

bool LatticeNetDecode::decode_box(
    InferBoxRaw const& box_raw, int const& class_number, int const& class_base,
    ObjDecodeParameter const& parameters,
    FixedArray<BoundingBox, kMaxBoundingBoxNum>& bounding_boxes,
    std::string& err) noexcept {
  auto const& boxInfoNumber{parameters.bbox_info_number};
  auto const& netRange{parameters.net_range};
  auto const& voxelSize{parameters.voxel_size};
  auto const& downSize{parameters.down_size};

  int netLength =
      int((netRange[3] - netRange[0]) / voxelSize[0] / downSize + 0.5);

  int netWidth =
      int((netRange[4] - netRange[1]) / voxelSize[1] / downSize + 0.5);

  std::array<float, 10> classThreshold{0.4, 0.4, 0.4, 0.4, 0.4,
                                       0.4, 0.4, 0.4, 0.4, 0.4};

  // decode
  auto nbDecode{decode_chw_center(
      box_raw.cls.data(), box_raw.reg.data(), box_raw.height.data(),
      box_raw.dim.data(), box_raw.rot.data(), box_raw.vel.data(),
      classThreshold.data(), obj_scores_ptr->data(), obj_boxes_ptr->data(),
      obj_labels_ptr->data(),
      static_cast<int>(LatticeNetDecode::kMaxResultSize), 0, netLength,
      netWidth, class_number, class_base, boxInfoNumber, downSize,
      voxelSize.data(), netRange.data())};
  if (nbDecode == -1) {
    err = "decode object error";
    return -1;
  }

  // nms
  obj_nms_keep_ptr->fill(true);
  nms_cpu(obj_boxes_ptr->data(), obj_scores_ptr->data(),
          obj_nms_keep_ptr->data(), boxInfoNumber, parameters.nms_iou_threshold,
          nbDecode);

  for (int i = 0; i < nbDecode; i++) {
    if (obj_nms_keep_ptr->at(i)) {
      auto& box{bounding_boxes.GetNext()};
      box.score = obj_scores_ptr->at(i);
      box.label = static_cast<std::uint16_t>(obj_labels_ptr->at(i));
      box.center.x = obj_boxes_ptr->at(i * boxInfoNumber + 0);
      box.center.y = obj_boxes_ptr->at(i * boxInfoNumber + 1);
      box.center.z = obj_boxes_ptr->at(i * boxInfoNumber + 2);
      box.length = obj_boxes_ptr->at(i * boxInfoNumber + 3);
      box.width = obj_boxes_ptr->at(i * boxInfoNumber + 4);
      box.height = obj_boxes_ptr->at(i * boxInfoNumber + 5);
      box.r = obj_boxes_ptr->at(i * boxInfoNumber + 6);
      box.heading = obj_boxes_ptr->at(i * boxInfoNumber + 6);
      // box.vx = out_boxes.at(i * boxInfoNumber + 7);
      // box.vy = out_boxes.at(i * boxInfoNumber + 8);
      // box.ax = out_boxes.at(i * boxInfoNumber + 9);
      // box.ay = out_boxes.at(i * boxInfoNumber + 10);

      gen_polygon(box, box.polygon);
      if (!bounding_boxes.IsFull()) {
        bounding_boxes.Add();
      }
    }
  }
  return true;
}

bool LatticeNetDecode::DecodeObject(
    InferOutputDataRaw const& dataRaw, ObjDecodeParameter const& parameters,
    FixedArray<BoundingBox, kMaxBoundingBoxNum>& boundingBoxes,
    std::string& err) noexcept {
  std::string errorMsg;
  std::array<int, 6> clsNumberArray{1, 2, 2, 1, 2, 2};
  int classBase = 0;

  boundingBoxes.Clear();

  // box_0
  classBase =
      std::accumulate(clsNumberArray.begin(), clsNumberArray.begin() + 0, 0);
  bool bRet{decode_box(dataRaw.box_0, clsNumberArray[0], classBase, parameters,
                       boundingBoxes, errorMsg)};
  if (bRet != true) {
    err = "decode box_0 error, " + errorMsg;
    return false;
  }

  // box_1
  classBase =
      std::accumulate(clsNumberArray.begin(), clsNumberArray.begin() + 1, 0);
  bRet = decode_box(dataRaw.box_1, clsNumberArray[1], classBase, parameters,
                    boundingBoxes, errorMsg);
  if (bRet != true) {
    err = "decode box_1 error, " + errorMsg;
    return false;
  }

  // box_2
  classBase =
      std::accumulate(clsNumberArray.begin(), clsNumberArray.begin() + 2, 0);
  bRet = decode_box(dataRaw.box_2, clsNumberArray[2], classBase, parameters,
                    boundingBoxes, errorMsg);
  if (bRet != true) {
    err = "decode box_2 error, " + errorMsg;
    return false;
  }

  // box_3
  classBase =
      std::accumulate(clsNumberArray.begin(), clsNumberArray.begin() + 3, 0);
  bRet = decode_box(dataRaw.box_3, clsNumberArray[3], classBase, parameters,
                    boundingBoxes, errorMsg);
  if (bRet != true) {
    err = "decode box_3 error, " + errorMsg;
    return false;
  }

  // box_4
  classBase =
      std::accumulate(clsNumberArray.begin(), clsNumberArray.begin() + 4, 0);
  bRet = decode_box(dataRaw.box_4, clsNumberArray[4], classBase, parameters,
                    boundingBoxes, errorMsg);
  if (bRet != true) {
    err = "decode box_3 error, " + errorMsg;
    return false;
  }

  // box_5
  classBase =
      std::accumulate(clsNumberArray.begin(), clsNumberArray.begin() + 5, 0);
  bRet = decode_box(dataRaw.box_5, clsNumberArray[5], classBase, parameters,
                    boundingBoxes, errorMsg);
  if (bRet != true) {
    err = "decode box_5 error, " + errorMsg;
    return false;
  }

  return true;
}

void LatticeNetDecode::gen_polygon(BoundingBox const& bbox,
                                   FixedArray<Point2f, 4>& polygon) noexcept {
  Eigen::Vector2f box_centor, dirt;
  Eigen::Vector2f dirt_length;
  Eigen::Vector2f dirt_width;
  box_centor[0] = bbox.center.x;
  box_centor[1] = bbox.center.y;
  dirt[0] = cos(bbox.heading);
  dirt[1] = sin(bbox.heading);
  dirt_length[0] = dirt[0];
  dirt_length[1] = dirt[1];
  dirt_width[0] = -dirt[1];
  dirt_width[1] = dirt[0];

  polygon.Clear();
  Eigen::Vector2f p{box_centor +
                    (bbox.length * dirt_length + bbox.width * dirt_width) / 2};
  auto& point0{polygon.GetNext()};
  point0.x = p[0];
  point0.y = p[1];
  polygon.Add();

  p = box_centor + (bbox.length * dirt_length - bbox.width * dirt_width) / 2;
  auto& point1{polygon.GetNext()};
  point1.x = p[0];
  point1.y = p[1];
  polygon.Add();

  p = box_centor - (bbox.length * dirt_length + bbox.width * dirt_width) / 2;
  auto& point2{polygon.GetNext()};
  point2.x = p[0];
  point2.y = p[1];
  polygon.Add();

  p = box_centor - (bbox.length * dirt_length - bbox.width * dirt_width) / 2;
  auto& point3{polygon.GetNext()};
  point3.x = p[0];
  point3.y = p[1];
  polygon.Add();
}

bool LatticeNetDecode::DecodeMap(InferOutputDataRaw const& dataRaw,
                                 MapDecodeParameter const& parameters, Map& map,
                                 std::string& err) noexcept {
  auto const& threshold{parameters.class_threshold};
  auto const& netRange{parameters.net_range};
  auto const& voxelSize{parameters.voxel_size};
  auto const& downSize{parameters.down_size};
  auto const& distanceThreshold{parameters.distance_threshold};
  auto const& mapInfoSize{parameters.map_info_number};

  int netLength =
      int((netRange[3] - netRange[0]) / voxelSize[0] / downSize + 0.5);

  int netWidth =
      int((netRange[4] - netRange[1]) / voxelSize[1] / downSize + 0.5);

  auto voxelSizeOut{voxelSize};
  voxelSizeOut[0] = voxelSizeOut[0] * downSize;
  voxelSizeOut[1] = voxelSizeOut[1] * downSize;

  // road edge
  auto rSize{decode_lane_double_chw(
      dataRaw.map_0.cls.data(), dataRaw.map_0.reg.data(), threshold[0],
      netLength, netWidth, 2, voxelSizeOut.data(), netRange.data(),
      distanceThreshold, kMaxResultSize, mapInfoSize, map_pts_ptr_->data())};
  if (rSize == -1) {
    err = "decode map 0 error";
    return false;
  }

  map.road_edge.Clear();
  std::unordered_map<std::string, RoadEdgeLine*> roadEdgeLineDict;
  for (auto i = 0; i < rSize; i++) {
    auto const& id{std::to_string(
        static_cast<int32_t>(map_pts_ptr_->at(i * mapInfoSize + 0)))};
    if (roadEdgeLineDict.find(id) == roadEdgeLineDict.end() &&
        !map.road_edge.IsFull()) {
      auto& roadEdgeline{map.road_edge.GetNext()};
      roadEdgeline.points.Clear();
      roadEdgeline.id = id;
      map.road_edge.Add();
      roadEdgeLineDict.emplace(id, &roadEdgeline);
    }
    auto roadEdgeLinePtr{roadEdgeLineDict[id]};
    if (roadEdgeLinePtr != nullptr && !roadEdgeLinePtr->points.IsFull()) {
      auto& point{roadEdgeLinePtr->points.GetNext()};
      point.x = map_pts_ptr_->at(i * mapInfoSize + 1);
      point.y = map_pts_ptr_->at(i * mapInfoSize + 2);
      point.z = map_pts_ptr_->at(i * mapInfoSize + 3);
      roadEdgeLinePtr->points.Add();
    }
  }

  // lane divider
  rSize = decode_lane_double_chw(
      dataRaw.map_1.cls.data(), dataRaw.map_1.reg.data(), threshold[1],
      netLength, netWidth, 2, voxelSizeOut.data(), netRange.data(),
      distanceThreshold, kMaxResultSize, mapInfoSize, map_pts_ptr_->data());
  if (rSize == -1) {
    err = "decode map 1 error";
    return false;
  }

  map.lane_divider.Clear();
  std::unordered_map<std::string, LaneDividerLine*> laneDividerLineDict;
  for (auto i = 0; i < rSize; i++) {
    auto const& id{std::to_string(
        static_cast<int32_t>(map_pts_ptr_->at(i * mapInfoSize + 0)))};
    if (laneDividerLineDict.find(id) == laneDividerLineDict.end() &&
        !map.lane_divider.IsFull()) {
      auto& lineDividerline{map.lane_divider.GetNext()};
      lineDividerline.points.Clear();
      lineDividerline.id = id;
      map.lane_divider.Add();
      laneDividerLineDict.emplace(id, &lineDividerline);
    }
    auto lineDividerLinePtr{laneDividerLineDict[id]};
    if (lineDividerLinePtr != nullptr && !lineDividerLinePtr->points.IsFull()) {
      auto& point{lineDividerLinePtr->points.GetNext()};
      point.x = map_pts_ptr_->at(i * mapInfoSize + 1);
      point.y = map_pts_ptr_->at(i * mapInfoSize + 2);
      point.z = map_pts_ptr_->at(i * mapInfoSize + 3);
      lineDividerLinePtr->points.Add();
    }
  }

  // road marking
  rSize = decode_lane_double_chw(
      dataRaw.map_2.cls.data(), dataRaw.map_2.reg.data(), threshold[2],
      netLength, netWidth, 2, voxelSizeOut.data(), netRange.data(),
      distanceThreshold, kMaxResultSize, mapInfoSize, map_pts_ptr_->data());
  if (rSize == -1) {
    err = "decode map 2 error";
    return false;
  }

  map.road_marking.Clear();
  std::unordered_map<std::string, RoadMarkingLine*> roadMarkingLineDict;
  for (auto i = 0; i < rSize; i++) {
    auto const& id{std::to_string(
        static_cast<int32_t>(map_pts_ptr_->at(i * mapInfoSize + 0)))};
    if (roadMarkingLineDict.find(id) == roadMarkingLineDict.end() &&
        !map.road_marking.IsFull()) {
      auto& roadMarkingline{map.road_marking.GetNext()};
      roadMarkingline.polygons.Clear();
      roadMarkingline.categorys.Clear();
      roadMarkingline.id = id;
      map.road_marking.Add();
      roadMarkingLineDict.emplace(id, &roadMarkingline);

      auto& category{roadMarkingline.categorys.GetNext()};
      category = "crosswalk";
      roadMarkingline.categorys.Add();
    }
    auto roadMarkingLinePtr{roadMarkingLineDict[id]};
    if (roadMarkingLinePtr != nullptr &&
        !roadMarkingLinePtr->polygons.IsFull()) {
      auto& polygon{roadMarkingLinePtr->polygons.GetNext()};
      polygon.x = map_pts_ptr_->at(i * mapInfoSize + 1);
      polygon.y = map_pts_ptr_->at(i * mapInfoSize + 2);
      polygon.z = map_pts_ptr_->at(i * mapInfoSize + 3);
      roadMarkingLinePtr->polygons.Add();
    }
  }
  return true;
}